package main

import (
	"io"
	"log"
	"net"
	"sync"
)

const (
	LocalAddr  = "127.0.0.1:6891"   // 本地监听地址
	RemoteAddr = "10.225.9.94:6891" // 目标服务器地址
)

func main() {
	// 监听本地端口，用于接收客户端连接
	listener, err := net.Listen("tcp", ":6891")
	if err != nil {
		log.Fatalf("Failed to start proxy server: %v", err)
	}
	defer listener.Close()

	log.Println("Proxy server is running on :6891")

	for {
		// 接受客户端连接
		clientConn, err := listener.Accept()
		if err != nil {
			log.Printf("Failed to accept client connection: %v", err)
			continue
		}

		// 处理每个客户端连接
		go handleClient(clientConn)
	}
}

func handleClient(clientConn net.Conn) {
	defer clientConn.Close()

	// 连接到目标服务器
	targetConn, err := net.Dial("tcp", RemoteAddr) // 假设目标服务器运行在 localhost:9090
	if err != nil {
		log.Printf("Failed to connect to target server: %v", err)
		return
	}
	defer targetConn.Close()

	log.Println("Connected to target server")

	// 使用 WaitGroup 确保两个方向的数据传输完成后再退出
	var wg sync.WaitGroup
	wg.Add(2)

	// 客户端 -> 目标服务器
	go func() {
		defer wg.Done()
		io.Copy(targetConn, clientConn)
		clientConn.Close() // 当客户端到目标服务器的数据流结束时，关闭客户端连接
	}()

	// 目标服务器 -> 客户端
	go func() {
		defer wg.Done()
		io.Copy(clientConn, targetConn)
		targetConn.Close() // 当目标服务器到客户端的数据流结束时，关闭目标服务器连接
	}()

	// 等待两个方向的数据传输完成
	wg.Wait()
	log.Println("Connection closed")
}
